# 数据集验证，需要加载的包库
import os
from PIL import Image
from matplotlib import pyplot as plt
import xml
import cv2 as cv
import numpy as np
import json
import math
from scipy.spatial import distance
from collections import OrderedDict
import time
import pandas as pds
from torch.utils.data import DataLoader
from torchvision import models
import torch
import torch.nn as nn
from utilour.datasetLoader import MyDataset
import torch.optim as optimer
from logdir import loginfo
from utilour.tool import calaccurary,visualize_cam
from torchsummary import summary
from utilour.gradcadpp import GradCAM,GradCAMpp
import torchvision
from torchvision.utils import make_grid, save_image
import shutil

# 清理文件
if os.path.exists('./outputs'):
    shutil.rmtree(os.path.join('.','outputs'))
os.mkdir(os.path.join('.','outputs'))



# 加载数据
csvpath="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/dataset/newcardata/train_test_vail.csv" # 文件路径
traindataset=MyDataset(csvpath,"","train",errload=False)
#vaildataset=MyDataset(csvpath,"vail")
#testdataset=MyDataset(csvpath,"test")
trainLoader=DataLoader(dataset=traindataset,batch_size=2,num_workers=8)
#vailLoader=DataLoader(dataset=vaildataset,batch_size=2,num_workers=8)
#testLoader=DataLoader(dataset=testdataset,batch_size=2,num_workers=8)


# 加载模型--修改模型
resnet50cpk=torch.load("/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/arckive/modeltrain12_24_23_12/modelcpk/model_minvl_13_8068.pkl")
#summary(resnet50cpk,input_size=(3,217,217))

# 声明gradcam--gradcampp  _cam  _campp
gradcam=GradCAM({"type":"resnet50","model":resnet50cpk,"target_layer":resnet50cpk._modules["features"]._modules["norm5"]})
gradcampp=GradCAMpp({"type":"resnet50","model":resnet50cpk,"target_layer":resnet50cpk._modules["features"]._modules["norm5"]})


totalnum=0
errornum=0

def outputImage(images,cvname,chdrs,sig="True",ispp=False,clsnum=5):
    nrows=10
    images = make_grid(torch.cat(images, 0), nrow=nrows)        

    # 维度调整，准备输出图片
    print(images.shape)
    imgs=images.numpy()*255
    imgs=np.transpose(imgs,[1,2,0])

    # 图片输出
    output_dir =os.path.join(".", 'outputs')
    os.makedirs(output_dir, exist_ok=True)
    if not ispp:
        imgs=cv.cvtColor(imgs.astype(np.uint8),cv.COLOR_BGR2RGB)
    print("image output path:{}".format(os.path.join(output_dir,"sers_{}_{}_{}".format(str(chdrs),sig,cvname))))
    cv.imwrite(os.path.join(output_dir,"sers_{}_{}_{}".format(str(chdrs),sig,cvname)),imgs.astype(np.uint8))

# 使用程序计算CAM 
def renderCAMHotMap(camprogram,dataloader,cvname='test_cv_cam.jpg',PILname="test_PIL_cam.jpg",ispp=False,errordict={},dataloadertype='train',sig=True):
    font=cv.FONT_HERSHEY_SIMPLEX#使用默认字体
    images=[]
    # 设置一个解算图片
    chdrs=0
    # 求解grad_CAM 的计算情况
    outsig="True" if sig else "False"
    for step,(data,label,imgpath) in  enumerate(dataloader):
        data=data.cuda()
        b,c,h,w=data.size()
        label=label.cuda()
        #tlabels=resnet50cpk(data)
        #saliency_maps,tlabel=gradcam.calGradCAM(data)
        calres=camprogram.calGradhot(data,clsn=5)
        labellist=label.cpu().numpy()
        imglist=data.cpu()

        for b in range(len(calres)):
            global totalnum
            totalnum=totalnum+1
            #原始图片
            img=imglist[b,:,:,:]
            tempimgs=[]
            tempinfo=calres[b] # 获取当前的类别计算情况
            tl=labellist[b]
            reslabel=tempinfo['label'].detach().cpu().numpy()
            gl=np.argmax(reslabel)
            if sig: # 记录正确
                if tl!=gl:
                    continue
            else:
                if tl==gl:
                    continue
            # 记录错误列表信息
            global errornum
            errornum=errornum+1
            csvfile['id'].append(len(csvfile.keys()))
            csvfile['imgname'].append(os.path.basename(imgpath[b]))
            csvfile['cls'].append(str(tl))
            csvfile['sig'].append(str(gl))
            csvfile['trval'].append(dataloadertype)

            # 图片的信息
            imginfo_cam=np.zeros((h,w,c),np.uint8) # 新建图像，注意一定要是uint8
            imginfo_cam=cv.putText(imginfo_cam,'{}'.format(os.path.basename(imgpath[b])),(0,50),font,0.6,(255,255,255),2)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细
            imginfo_cam=cv.putText(imginfo_cam,'S:{}\nT:{}'.format(tl,gl),(0,100),font,0.8,(255,255,255),2)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细
            imginfo_cam=imginfo_cam/255.0
            imginfo_cam=torch.from_numpy(np.transpose(imginfo_cam,[2,0,1]))
            tempimgs.append(imginfo_cam)


            # 标签信息
            labelinfo_cam=np.zeros((h,w,c),np.uint8) # 新建图像，注意一定要是uint8
            for k in range(len(tempinfo["saliency"])):
                labelinfo_cam=cv.putText(labelinfo_cam,'{}:{}'.format(k,'%.3f'%reslabel[k]),(0,40*(k+1)),font,0.9,(255,255,255),2)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细
            labelinfo_cam=labelinfo_cam/255.0
            labelinfo_cam=torch.from_numpy(np.transpose(labelinfo_cam,[2,0,1]))
            tempimgs.append(labelinfo_cam)

            # 类别渲染情况
            for k in range(len(tempinfo["saliency"])):
                saliencyinfo=tempinfo["saliency"][k]
                # saliency_map 
                saliency_map_cam=saliencyinfo["data"]
                heatmap, result = visualize_cam(saliency_map_cam.detach().cpu().numpy(), img)
                #tempimgs.append(img.squeeze().cpu())
                #tempimgs.append(heatmap)
                tempimgs.append(result)

            # 热力图相关信息
            hotmapmin=np.zeros((h,w,c),np.uint8)
            for k in range(len(tempinfo["saliency"])):
                hotmapmin=cv.putText(hotmapmin,'{}:{}'.format('min_{}'.format(k),'%.7f'%tempinfo["saliency"][k]["min"]),(0,30*(k+1)),font,0.8,(255,255,255),2)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细
            hotmapmin=hotmapmin/255.0
            hotmapmin=torch.from_numpy(np.transpose(hotmapmin,[2,0,1]))
            tempimgs.append(hotmapmin)

            hotmapmax=np.zeros((h,w,c),np.uint8)
            for k in range(len(tempinfo["saliency"])):
                hotmapmax=cv.putText(hotmapmax,'{}:{}'.format('max_{}'.format(k),'%.7f'%tempinfo["saliency"][k]["max"]),(0,30*(k+1)),font,0.8,(255,255,255),2)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细
            hotmapmax=hotmapmax/255.0
            hotmapmax=torch.from_numpy(np.transpose(hotmapmax,[2,0,1]))
            tempimgs.append(hotmapmax)
            # 添加原始图片
            tempimgs.append(img.squeeze().cpu())
            
            images.append(torch.stack(tempimgs, 0))
        if len(images)>=30:
            print(len(images))
            outputImage(images,cvname,chdrs,sig=outsig,ispp=ispp,clsnum=5)
            chdrs=chdrs+1
            images=[]
    print(len(images))
    outputImage(images,cvname,chdrs,sig=outsig,ispp=ispp,clsnum=5)
    return errordict # 返回了错误dict
    #d=Image.fromarray(imgs.astype(np.uint8))
    #d.save(os.path.join(output_dir,PILname ))


csvfile={'id':[],'imgname':[],'cls':[],'sig':[],'trval':[]} # 创建错误记录列表

# 使用test,分别计算 grad_cam,grad_campp
#renderCAMHotMap(gradcam,testLoader,cvname="test_cv_cam.jpg",PILname="test_PIL_cam.jpg")
#print("test,cam,over acc:{}".format((totalnum-errornum)/totalnum))
#renderCAMHotMap(gradcampp,testLoader,cvname="test_cv_campp.jpg",PILname="test_PIL_campp.jpg",errordict=csvfile,dataloadertype='test')
#print("test,campp,over acc:{}".format((totalnum-errornum)/totalnum))


# 使用train,分别计算 grad_cam,grad_campp
renderCAMHotMap(gradcam,trainLoader,cvname="train_cv_cam.jpg",PILname="train_PIL_cam.jpg",sig=True)
print("train,cam,over acc:{}".format((totalnum-errornum)/totalnum))
renderCAMHotMap(gradcampp,trainLoader,cvname="train_cv_campp.jpg",PILname="train_PIL_campp.jpg",errordict=csvfile,dataloadertype='train',sig=True)
print("train,campp,over acc:{}".format((totalnum-errornum)/totalnum))

renderCAMHotMap(gradcam,trainLoader,cvname="train_cv_cam.jpg",PILname="train_PIL_cam.jpg",sig=False)
print("train,cam,over acc:{}".format((totalnum-errornum)/totalnum))
renderCAMHotMap(gradcampp,trainLoader,cvname="train_cv_campp.jpg",PILname="train_PIL_campp.jpg",errordict=csvfile,dataloadertype='train',sig=False)
print("train,campp,over acc:{}".format((totalnum-errornum)/totalnum))

# 使用train,分别计算 grad_cam,grad_campp
#renderCAMHotMap(gradcam,vailLoader,cvname="vail_cv_cam.jpg",PILname="vail_PIL_cam.jpg")
#print("vail,cam,over acc:{}".format((totalnum-errornum)/totalnum))
#renderCAMHotMap(gradcampp,vailLoader,cvname="vail_cv_campp.jpg",PILname="vail_PIL_campp.jpg",errordict=csvfile,dataloadertype='vail')
#print("vail,campp,over acc:{}".format((totalnum-errornum)/totalnum))

# 创建文件，记录错误文件
csvobj=pds.DataFrame(csvfile)
csvobj.to_csv( os.path.join('.','outputs','csvobj.csv'),encoding='utf-8')

# 压缩文件
localtime = time.localtime(time.time())
shelltext="zip  -r  ./testGram_{}_{}.zip  ./outputs".format(localtime.tm_mon,localtime.tm_mday)
print('文件处理完成，正在进行文件压缩命令 :{} ，请勿关闭软件。。。。'.format(shelltext))
os.system(shelltext)

